#include "BackupFSFileNodes.h"
#include "logger/BackupFSLogger.h"
#include "BackupFSSql.h"
#include <sddl.h>
#include <thread>
#include <chrono>

namespace BackupFS
{
	BackupFSFileNodes::BackupFSFileNodes()
	{
		_StopFS = false;
		std::thread t(&BackupFSFileNodes::FlushFS, this);
		t.detach();
	}
	BackupFSFileNodes::~BackupFSFileNodes()
	{
		Stop();
		_DirectoryPaths.clear();
		_FileNodes.clear();
		Logger->Info("Release all the resources from nodes");
	}
	void BackupFSFileNodes::Stop()
	{
		if (_StopFS)
		{
			return;
		}
		_StopFS = true;
		std::scoped_lock lock(_StopFSMutex);
		BackupFSSql::Stop(BackupFSCore::GetInstance()->GetDb());
		BackupFSCore::GetInstance()->StopFS();
	}
	/*
	Flush the meta data
	*/
	void BackupFSFileNodes::FlushFS()
	{
		std::scoped_lock lock(_StopFSMutex);
		BackupFSSql::Begin(BackupFSCore::GetInstance()->GetDb());
		std::chrono::system_clock::time_point begin = std::chrono::system_clock::now();
		while (!_StopFS)
		{
			std::chrono::system_clock::time_point now = std::chrono::system_clock::now();
			std::chrono::duration<double> elapsed_seconds = now - begin;
			int flush_val = _FlushNow.load();
			if (elapsed_seconds.count() >= _FlushFrequency|| flush_val==1)
			{
				if (!BackupFSSql::Commit(BackupFSCore::GetInstance()->GetDb()))
				{
					break;
				}
				if (!BackupFSSql::Begin(BackupFSCore::GetInstance()->GetDb()))
				{
					break;
				}
				begin = std::chrono::system_clock::now();
				if (flush_val == 1)
				{
					_FlushNow.store(0);
				}
			}
			else
			{
				/*
				When cache is full, _FlushNow will be set to 1
				Need to commit immediately
				*/
				for (int i = 0; i < 10; i++)
				{
					int flush_val = _FlushNow.load();
					if (flush_val == 1||_StopFS)
					{
						break;
					}
					std::this_thread::sleep_for(std::chrono::milliseconds(100));
				}
			}
		}
		Logger->Info("Stop FS from Nodes");
	}
	void BackupFSFileNodes::AddRoot(bool add_to_db)
	{
		_AddRoot(add_to_db);
	}
	/*
	Add the root \ node
	It's the root path such as x:\ or c:\mountpoint
	*/
	void BackupFSFileNodes::_AddRoot(bool add_to_db)
	{
		WCHAR buffer[1024] = { 0 };
		WCHAR final_buffer[2048] = { 0 };
		PTOKEN_USER user_token = nullptr;
		PTOKEN_GROUPS groups_token = nullptr;
		HANDLE token_handle;
		LPTSTR user_sid_str = nullptr;
		LPTSTR group_sid_str = nullptr;
		PSECURITY_DESCRIPTOR security_descriptor = nullptr;
		do
		{
			// Build default root filenode SecurityDescriptor
			if (!OpenProcessToken(GetCurrentProcess(), TOKEN_READ, &token_handle))
			{
				break;
			}
			DWORD return_length;
			if (!GetTokenInformation(token_handle, TokenUser, buffer, sizeof(buffer), &return_length))
			{
				CloseHandle(token_handle);
				break;
			}
			user_token = (PTOKEN_USER)buffer;
			if (!ConvertSidToStringSid(user_token->User.Sid, &user_sid_str))
			{
				CloseHandle(token_handle);
				break;
			}
			if (!GetTokenInformation(token_handle, TokenGroups, buffer, sizeof(buffer), &return_length))
			{
				CloseHandle(token_handle);
				break;
			}
			groups_token = (PTOKEN_GROUPS)buffer;
			if (groups_token->GroupCount > 0)
			{
				if (!ConvertSidToStringSid(groups_token->Groups[0].Sid, &group_sid_str))
				{
					CloseHandle(token_handle);
					break;
				}
				swprintf_s(buffer, 1024, L"O:%lsG:%ls", user_sid_str, group_sid_str);
			}
			else
			{
				swprintf_s(buffer, 1024, L"O:%ls", user_sid_str);
			}
			CloseHandle(token_handle);
			swprintf_s(final_buffer, 2048, L"%lsD:PAI(A;OICI;FA;;;AU)", buffer);
			ULONG size = 0;
			if (!ConvertStringSecurityDescriptorToSecurityDescriptor(
				final_buffer, SDDL_REVISION_1, &security_descriptor, &size))
			{
				break;
			}

			auto fileNode = std::make_shared<BackupFSFileNode>(L"\\", true, FILE_ATTRIBUTE_DIRECTORY, nullptr);
			fileNode->ParentId = 0;
			fileNode->NodeId = 1;
			fileNode->Security.SetDescriptor(security_descriptor);
			std::wstring root_path = L"\\";
			_FileNodes[root_path] = fileNode;
			_DirectoryPaths.emplace(root_path, std::set<std::shared_ptr<BackupFSFileNode>>());
			if (add_to_db)
			{
				_InsertToDb(fileNode);
			}
			BackupFSSql::GetAllChildren(BackupFSCore::GetInstance()->GetDb(), root_path, _DirectoryPaths.find(root_path)->second);
			for (auto &node : _DirectoryPaths.find(root_path)->second)
			{
				_FileNodes[node->GetFileName()] = node;
			}
		} while (0);

		if (user_sid_str)
		{
			LocalFree(user_sid_str);
		}
		if (group_sid_str)
		{
			LocalFree(group_sid_str);
		}
		if (security_descriptor)
		{
			LocalFree(security_descriptor);
		}
	}
	void BackupFSFileNodes::_ReleaseCache()
	{
		
		/*
		Currently use a simple method,remove all the cache and submit a commit request
		Can use LRU method in future
		*/
		_DirectoryPaths.clear();
		_FileNodes.clear();
		//commit the db
		_FlushNow.store(1);
		for (int i = 0; i < 3000; i++)
		{
			int flush_val = _FlushNow.load();
			if (flush_val == 0)
			{
				break;
			}
			std::this_thread::sleep_for(std::chrono::milliseconds(1));
		}

	}
	void BackupFSFileNodes::_UpdateNodeLRU(const std::wstring& filename)
	{
		/*if (filename == L"\\")
		{
			return;
		}
		_NodesCache.Update(filename, time(NULL));*/
	}
	/*
	Add a new node to the file system

	*/
	NTSTATUS BackupFSFileNodes::Add(const std::shared_ptr<BackupFSFileNode>& f,
		std::optional<std::pair<std::wstring, std::wstring>> stream_names)
	{
		std::scoped_lock lock(_FilesNodesMutex);
		if (_NodesCacheNumber < _FileNodes.size())
		{
			// release the cache
			_ReleaseCache();
		}
		if (f->FileIndex == 0)  // previous init
			f->FileIndex = _FsFileIndexCount++;
		const auto filename = f->GetFileName();
		const auto parent_path = BackupFSUtils::GetParentPath(filename);

		// Does target folder exist
		if (!_DirectoryPaths.count(parent_path))
		{
			ListFolder(parent_path);
			if (!_DirectoryPaths.count(parent_path))
			{
				Logger->Error(W2S(parent_path), " doesn't exist");
				return STATUS_OBJECT_PATH_NOT_FOUND;
			}
		}
		LONGLONG parent_node_id = 0;
		const auto parent_node = _FileNodes[parent_path];
		if (parent_node)
		{
			parent_node_id = parent_node->NodeId;
		}
		if (!stream_names.has_value())
		{
			stream_names = BackupFSUtils::GetStreamNames(filename);
		}
		//Check the ADS
		if (!stream_names.value().second.empty())
		{
			auto& stream_names_value = stream_names.value();
			auto main_stream_name = BackupFSUtils::GetFileNameStreamLess(filename, stream_names_value);
			auto main_f = Find(main_stream_name);
			if (!main_f)
			{
				return STATUS_OBJECT_PATH_NOT_FOUND;
			}
			main_f->AddStream(f);
			f->MainStream = main_f;
			f->FileIndex = main_f->FileIndex;
		}

		// If we have a folder, we add it to our directoryPaths
		if (f->IsDirectory && !_DirectoryPaths.count(filename))
		{
			_DirectoryPaths.emplace(filename, std::set<std::shared_ptr<BackupFSFileNode>>());
			_UpdateNodeLRU(filename);
			Logger->Debug("add path ", W2S(filename));
		}

		// Add our file to the fileNodes and directoryPaths
		auto previous_f = _FileNodes[filename];
		_FileNodes[filename] = f;
		_DirectoryPaths[parent_path].insert(f);
		
		f->ParentId = parent_node_id;
		f->NodeId = BackupFSCore::FsInodeNumber++;
		_InsertToDb(f);
		Logger->Debug("add node ", W2S(filename));
		if (previous_f)
		{
			Logger->Info("erase the previous  ", W2S(filename));
			_DirectoryPaths[parent_path].erase(previous_f);
			BackupFSSql::RemoveFile(BackupFSCore::GetInstance()->GetDb(), previous_f);
		}
		_UpdateNodeLRU(parent_path);
		return STATUS_SUCCESS;
	}
	bool BackupFSFileNodes::_InsertToDb(std::shared_ptr<BackupFSFileNode> node)
	{
		return BackupFSSql::AddFileMeta(BackupFSCore::GetInstance()->GetDb(), node);
	}
	std::shared_ptr<BackupFSFileNode> BackupFSFileNodes::Find(const std::wstring& filename) 
	{
		std::scoped_lock lock(_FilesNodesMutex);
		auto fileNode = _FileNodes.find(filename);
		if (fileNode != _FileNodes.end())
		{
			return fileNode->second;
		}

		std::wstring parent_path = BackupFSUtils::GetParentPath(filename);
		if (_DirectoryPaths.find(parent_path) == _DirectoryPaths.end())
		{
			ListFolder(parent_path);
			auto file = _FileNodes.find(filename);
			if (file != _FileNodes.end())
			{
				return file->second;
			}
		}
		return nullptr;
	}

	std::set<std::shared_ptr<BackupFSFileNode>> BackupFSFileNodes::ListFolder(
		const std::wstring& fileName) 
	{
		std::scoped_lock lock(_FilesNodesMutex);

		auto it = _DirectoryPaths.find(fileName);
		if (it != _DirectoryPaths.end())
		{
			return it->second;
		}
		auto file = BackupFSSql::GetFileNode(BackupFSCore::GetInstance()->GetDb(), fileName);
		if (file)
		{
			_DirectoryPaths.emplace(fileName, std::set<std::shared_ptr<BackupFSFileNode>>());
			_UpdateNodeLRU(fileName);
			//get all the children
			BackupFSSql::GetAllChildren(BackupFSCore::GetInstance()->GetDb(), fileName, _DirectoryPaths.find(fileName)->second);
			for (auto &node : _DirectoryPaths.find(fileName)->second)
			{
				_FileNodes[node->GetFileName()] = node;
			}
			_FileNodes[fileName] = file;
			return  _DirectoryPaths.find(fileName)->second;
		}
		return std::set<std::shared_ptr<BackupFSFileNode>>();
	}

	void BackupFSFileNodes::Remove(const std::wstring& filename) 
	{
		return Remove(Find(filename));
	}

	void BackupFSFileNodes::Remove(const std::shared_ptr<BackupFSFileNode>& f) 
	{
		if (!f)
		{
			return;
		}

		std::scoped_lock lock(_FilesNodesMutex);
		auto fileName = f->GetFileName();

		// Remove node from fileNodes and directoryPaths
		_FileNodes.erase(fileName);
		_DirectoryPaths[BackupFSUtils::GetParentPath(fileName)].erase(f);
		// Remove the data from db
		LONGLONG removed_size = ROUND_UP_PAGES(f->FileSize) * BACKUPFS_PAGE_SIZE;
		BackupFSSql::UpdateCurrentUsedSize(BackupFSCore::GetInstance()->GetDb(), BackupFSCore::GetInstance()->AddUsedSize(-removed_size));
		BackupFSSql::RemoveFile(BackupFSCore::GetInstance()->GetDb(), f);
		if (!f->IsDirectory)
		{
			BackupFSCore::GetInstance()->RemoveFileBlocks(f->GetDataBlocks());
		}
		// if it was a directory we need to remove it from directoryPaths
		if (f->IsDirectory) 
		{
			// but first we need to remove the directory content by looking recursively
			// into it
			auto files = ListFolder(fileName);
			for (const auto& file : files)
			{
				Remove(file);
			}
			_DirectoryPaths.erase(fileName);
			_NodesCache.Remove(fileName);
		}

		// Cleanup streams
		if (f->MainStream) 
		{
			// Is an alternate stream
			f->MainStream->RemoveStream(f);
		}
		else 
		{
			// Is a main stream
			// Remove possible alternate stream
			for (const auto& [stream_name, node] : f->GetStreams())
			{
				Remove(stream_name);
			}
		}
	}

	NTSTATUS BackupFSFileNodes::Move(const std::wstring& old_filename,
		const std::wstring& new_filename,
		BOOL replace_if_existing) 
	{
		auto f = Find(old_filename);
		auto new_f = Find(new_filename);

		if (!f)
		{
			return STATUS_OBJECT_NAME_NOT_FOUND;
		}

		// Cannot move to an existing destination without replace flag
		if (!replace_if_existing && new_f)
		{
			return STATUS_OBJECT_NAME_COLLISION;
		}

		// Cannot replace read only destination
		if (new_f && new_f->Attributes & FILE_ATTRIBUTE_READONLY)
		{
			return STATUS_ACCESS_DENIED;
		}

		// If destination exist - Cannot move directory or replace a directory
		if (new_f && (f->IsDirectory || new_f->IsDirectory))
		{
			return STATUS_ACCESS_DENIED;
		}

		auto newParentPath = BackupFSUtils::GetParentPath(new_filename);

		std::scoped_lock lock(_FilesNodesMutex);
		if (!_DirectoryPaths.count(newParentPath)) 
		{
			return STATUS_OBJECT_PATH_NOT_FOUND;
		}

		// Remove destination
		Remove(new_f);
		
		// Update current node with new data
		const auto fileName = f->GetFileName();
		auto oldParentPath = BackupFSUtils::GetParentPath(fileName);
		f->SetFileName(new_filename);
		//in the same folder
		_FileNodes[new_filename] = f;
		_FileNodes.erase(old_filename);
		LONGLONG parent_id = f->ParentId;
		if (oldParentPath != newParentPath)
		{
			//not in same folder
			_DirectoryPaths[newParentPath].insert(f);
			_DirectoryPaths[oldParentPath].erase(f);
			auto p = _FileNodes[newParentPath];
			if (p)
			{
				parent_id = p->NodeId;
			}
		}
		BackupFSSql::UpdateFileName(BackupFSCore::GetInstance()->GetDb(), f->NodeId, new_filename, parent_id);
		if (f->IsDirectory)
		{
			BackupFSSql::UpdateChildrenFileName(BackupFSCore::GetInstance()->GetDb(), old_filename, new_filename);
		}

		return STATUS_SUCCESS;
	}
}  // namespace BackupFS
